import { useEffect, useRef, useState, useCallback } from 'react';
import { TaskProgressWebSocketClient, TaskProgressInfo } from '../utils/TaskProgressWebSocketClient';

/**
 * 任务进度WebSocket Hook
 * 提供简化的API来订阅和管理任务进度更新
 */
export const useTaskProgress = (serverUrl?: string) => {
  const clientRef = useRef<TaskProgressWebSocketClient | null>(null);
  const [isConnected, setIsConnected] = useState(false);
  const [currentProgress, setCurrentProgress] = useState<TaskProgressInfo | null>(null);

  // 初始化WebSocket客户端
  useEffect(() => {
    clientRef.current = new TaskProgressWebSocketClient(serverUrl);
    
    // 设置连接状态回调
    clientRef.current.setConnectionStatusCallback(setIsConnected);
    
    // 连接到服务器
    clientRef.current.connect();

    // 清理函数
    return () => {
      if (clientRef.current) {
        clientRef.current.disconnect();
        clientRef.current = null;
      }
    };
  }, [serverUrl]);

  // 订阅任务进度
  const subscribeToTask = useCallback((taskId: number, onProgress?: (progress: TaskProgressInfo) => void) => {
    if (!clientRef.current) return;

    const handleProgress = (progress: TaskProgressInfo) => {
      setCurrentProgress(progress);
      onProgress?.(progress);
    };

    clientRef.current.subscribeToTaskProgress(taskId, handleProgress);
  }, []);

  // 取消订阅任务进度
  const unsubscribeFromTask = useCallback((taskId: number) => {
    if (!clientRef.current) return;
    clientRef.current.unsubscribeFromTaskProgress(taskId);
  }, []);

  // 查询任务进度
  const queryProgress = useCallback((taskId: number) => {
    if (!clientRef.current) return;
    clientRef.current.queryTaskProgress(taskId);
  }, []);

  // 发送ping
  const ping = useCallback(() => {
    if (!clientRef.current) return;
    clientRef.current.ping();
  }, []);

  return {
    isConnected,
    currentProgress,
    subscribeToTask,
    unsubscribeFromTask,
    queryProgress,
    ping
  };
};

/**
 * 特定任务的进度Hook
 * 自动管理单个任务的订阅和取消订阅
 */
export const useSpecificTaskProgress = (taskId: number | null, serverUrl?: string) => {
  const { isConnected, currentProgress, subscribeToTask, unsubscribeFromTask } = useTaskProgress(serverUrl);
  const [taskProgress, setTaskProgress] = useState<TaskProgressInfo | null>(null);

  useEffect(() => {
    if (!taskId || !isConnected) return;

    // 订阅任务进度
    subscribeToTask(taskId, setTaskProgress);

    // 清理：取消订阅
    return () => {
      unsubscribeFromTask(taskId);
    };
  }, [taskId, isConnected, subscribeToTask, unsubscribeFromTask]);

  return {
    isConnected,
    progress: taskProgress,
    taskId
  };
};